import cv2
import numpy as np
import os
from PIL import ImageGrab

def segment_chessboard(image, rows, cols):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
    
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    contours = sorted(contours, key=lambda c: cv2.boundingRect(c)[1])
    
    segmented = image.copy()
    
    chess_pieces = [[None for _ in range(cols)] for _ in range(rows)]
    
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        
        cv2.rectangle(segmented, (x, y), (x+w, y+h), (0, 0, 255), 2)
        
        row = int((y + h/2) // (image.shape[0] / rows)) + 1
        col = int((x + w/2) // (image.shape[1] / cols)) + 1
        
        chess_pieces[row-1][col-1] = image[y:y+h, x:x+w]
    
    for row in range(rows):
        for col in range(cols):
            cv2.imwrite(f"fenge/chess_piece_{row+1}_{col+1}.jpg", chess_pieces[row][col])
    
    return segmented

clipboard_image = np.array(ImageGrab.grabclipboard())

if clipboard_image is not None:
    rows = 16
    cols = 10

    segmented_chessboard = segment_chessboard(clipboard_image, rows, cols)

    #cv2.imshow("Segmented Chessboard", segmented_chessboard)
    #cv2.waitKey(0)
    #cv2.destroyAllWindows()
else:
    print("无法从剪贴板读取图像")

templates = []
for i in range(1, 10):
    template_path = os.path.join('muban', f'template_{i}.jpg')
    template = cv2.imread(template_path, cv2.IMREAD_GRAYSCALE)
    templates.append(template)

chess_board = [[0] * 10 for _ in range(16)]  # 创建一个16x10的棋盘二维数组

for x in range(16):  # 假设棋盘有16行
    for y in range(10):  # 假设棋盘有10列
        image_path = os.path.join('fenge', f'chess_piece_{x+1}_{y+1}.jpg')
        
        if not os.path.isfile(image_path):
            print(f'图像文件不存在：{image_path}')
            continue
        
        image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)

        if image is None:
            print(f'无法加载图像：{image_path}')
            continue

        best_match = None
        best_similarity = 0

        for j, template in enumerate(templates):
            res = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)

            if max_val > best_similarity:
                best_match = j + 1  # 棋子数字为模板的索引加1
                best_similarity = max_val

        chess_board[x][y] = best_match

with open('numbers.txt', 'w') as f:
    for row in chess_board:
        f.write(' '.join(str(num) for num in row) + '\n')

for row in chess_board:
    print(row)